package main

import (
    "context"
    "fmt"
    "log"
    "math"
    "math/rand"
    "os"
    "strconv"
    "sync"
    "time"

    "github.com/milvus-io/milvus-sdk-go/v2/client"
    "github.com/milvus-io/milvus-sdk-go/v2/entity"
)

var (
    nEntities, dim = 3000, 512
    //nEntities, dim = 7000000, 512
    showResult = false
)

const (
    milvusAddr     = `localhost:19530`
    collectionName = "hello_ascend_test001"

    msgFmt                         = "\n==== %s ====\n"
    idCol, randomCol, embeddingCol = "ID", "random", "embeddings"
    topK                           = 10

    insertLimit = 10000

    argvLen = 3
)

func Min(x, y int) int {
    if x < y {
        return x
    }
    return y
}

func main() {
    args := os.Args
    argNum := len(args)
    if argNum >= argvLen {
        var err error
        nEntities, err = strconv.Atoi(args[1])
        if err != nil {
            fmt.Printf("invalid parameter argv[1]: %s\n", err.Error())
        }
        dim, err = strconv.Atoi(args[2])
        if err != nil {
            fmt.Printf("invalid parameter argv[2]: %s\n", err.Error())
        }
        if argNum >= argvLen+1 {
            showRet, err := strconv.Atoi(args[3])
            if err != nil {
                fmt.Printf("invalid parameter argv[3]: %s\n", err.Error())
            } else {
                if showRet != 0 {
                    showResult = true
                }
            }
        }
    }

    ctx := context.Background()

    fmt.Printf(msgFmt, "start connecting to Milvus")
    c, err := client.NewClient(ctx, client.Config{
        Address: milvusAddr,
    })
    if err != nil {
        log.Fatalf("failed to connect to milvus, err: %v", err)
    }
    defer c.Close()

    // delete collection if exists
    has, err := c.HasCollection(ctx, collectionName)
    if err != nil {
        log.Fatalf("failed to check collection exists, err: %v", err)
    }
    if has {
        c.DropCollection(ctx, collectionName)
    }

    // create collection
    fmt.Printf(msgFmt, "create collection `hello_ascend`")
    schema := &entity.Schema{
        CollectionName: collectionName,
        Description:    "hello_ascend is the simplest demo to introduce the APIs",
        AutoID:         false,
        Fields: []*entity.Field{
            {
                Name:       idCol,
                DataType:   entity.FieldTypeInt64,
                PrimaryKey: true,
                AutoID:     false,
            },
            {
                Name:       randomCol,
                DataType:   entity.FieldTypeDouble,
                PrimaryKey: false,
                AutoID:     false,
            },
            {
                Name:     embeddingCol,
                DataType: entity.FieldTypeFloatVector,
                TypeParams: map[string]string{
                    entity.TypeParamDim: fmt.Sprintf("%d", dim),
                },
            },
        },
    }

    if err := c.CreateCollection(ctx, schema, 1); err != nil {
        log.Fatalf("create collection failed, err: %v", err)
    }
    //defer c.DropCollection(ctx, collectionName)

    // generate data
    fmt.Printf(msgFmt, "start generate random entities")
    rand.Seed(time.Now().UnixNano())
    begin := time.Now()
    idList, randomList, embeddingList := GenerateData(nEntities, 0)
    end := time.Now()
    searchLatency := float32(end.Sub(begin)) / float32(time.Millisecond)
    fmt.Printf("generate data time: %f\n", searchLatency)

    // insert data
    fmt.Printf(msgFmt, "start inserting random entities")
    begin = time.Now()
    InsertPaged(ctx, c, idList, randomList, embeddingList, "")
    if err := c.Flush(ctx, collectionName, false); err != nil {
        log.Fatalf("failed to flush data, err: %v", err)
    }
    end = time.Now()
    searchLatency = float32(end.Sub(begin)) / float32(time.Millisecond)
    fmt.Printf("insert time: %f\n", searchLatency)

    // build index
    fmt.Printf(msgFmt, "start creating index ASCEND_FLAT_INT8")
    begin = time.Now()
    idx, err := entity.NewIndexAscendFlatInt8(entity.IP)
    if err != nil {
        log.Fatalf("failed to create ASCEND_FLAT_INT8 index, err: %v", err)
    }
    if err := c.CreateIndex(ctx, collectionName, embeddingCol, idx, false); err != nil {
        log.Fatalf("failed to create index, err: %v", err)
    }
    end = time.Now()
    searchLatency = float32(end.Sub(begin)) / float32(time.Millisecond)
    fmt.Printf("creating index time: %f\n", searchLatency)

    // load collection
    fmt.Printf(msgFmt, "start loading collection")
    begin = time.Now()
    err = c.LoadCollection(ctx, collectionName, false)
    if err != nil {
        log.Fatalf("failed to load collection, err: %v", err)
    }
    end = time.Now()
    searchLatency = float32(end.Sub(begin)) / float32(time.Millisecond)
    fmt.Printf("loading collection time: %f\n", searchLatency)

    TestPerformance(ctx, c, embeddingList)
}

func TestPerformance(ctx context.Context, c client.Client, embeddingList [][]float32) {
    fmt.Printf(msgFmt, "start searching based on vector similarity")
    batchSize := []int{1, 2, 4, 8, 16, 32, 64, 128}
    var vec2search [][]entity.Vector
    sp, _ := entity.NewIndexFlatSearchParam()
    fmt.Println("warm up")
    warmUpTimes := 1
    for i := 0; i < len(batchSize); i++ {
        if batchSize[i] > nEntities {
            continue
        }

        var tempSearch []entity.Vector
        for j := 0; j < batchSize[i]; j++ {
            tempSearch = append(tempSearch, entity.FloatVector(embeddingList[len(embeddingList)-j-1]))
        }
        vec2search = append(vec2search, tempSearch)

        for j := 0; j < warmUpTimes; j++ {
            _, err := c.Search(ctx, collectionName, []string{}, "", []string{idCol, randomCol},
                vec2search[i], embeddingCol, entity.IP, topK, sp)
            if err != nil {
                log.Fatalf("failed to search collection, err: %v", err)
            }
        }
    }
    fmt.Println("start searching")
    for i := 0; i < len(batchSize); i++ {
        if batchSize[i] > nEntities {
            continue
        }
        fmt.Printf("test batchsize: %d\n", batchSize[i])

        loopTimes := 1
        begin := time.Now()
        for j := 0; j < loopTimes; j++ {
            sRet, err := c.Search(ctx, collectionName, []string{}, "", []string{idCol, randomCol},
                vec2search[i], embeddingCol, entity.IP, topK, sp)
            if err != nil {
                log.Fatalf("failed to search collection, err: %v", err)
            }
            if showResult {
                fmt.Println("results:")
                for _, res := range sRet {
                    printResult(&res)
                }
            }
        }
        end := time.Now()
        searchLatency := float32(end.Sub(begin)) / float32(time.Millisecond)
        fmt.Printf("base:%d, dim:%d, searchnum:%d, topk:%d, duration:%f, QPS:%f\n", nEntities, dim, batchSize[i],
            topK, searchLatency, 1000.0*float32(batchSize[i])*float32(loopTimes)/searchLatency)
    }

    fmt.Println("\nlast base vector:")
    for i := 0; i < dim; i++ {
        fmt.Printf("%f, ", embeddingList[len(embeddingList)-1][i])
    }
    fmt.Println()
}

func InsertPaged(ctx context.Context, c client.Client, idList []int64, randomList []float64,
    embeddingList [][]float32, partitionName string) {
    log.Printf(msgFmt, "start inserting random entities")
    total := len(idList)
    for i := 0; i < total; i += insertLimit {
        curNum := Min(insertLimit, total-i)
        idColData := entity.NewColumnInt64(idCol, idList[i:i+curNum])
        randomColData := entity.NewColumnDouble(randomCol, randomList[i:i+curNum])
        embeddingColData := entity.NewColumnFloatVector(embeddingCol, dim, embeddingList[i:i+curNum])

        log.Printf("insert %d:%d / %d\n", i, i+curNum, total)
        if _, err := c.Insert(ctx, collectionName, partitionName, idColData, randomColData, embeddingColData); err != nil {
            log.Fatalf("failed to insert random data into `hello_ascend`, err: %v", err)
        }
    }
}

// generate data
func GenerateData(total int, startID int) ([]int64, []float64, [][]float32) {
    fmt.Printf("generate data ntotal: %d, dim: %d\n", total, dim)
    rand.Seed(time.Now().UnixNano())
    idList, randomList := make([]int64, total, total), make([]float64, total, total)
    embeddingList := make([][]float32, total, total)
    var wg sync.WaitGroup
    for i := 0; i < total; i++ {
        wg.Add(1)
        go func(i int) {
            idList[i] = int64(i + startID)
            randomList[i] = rand.Float64()
            embeddingList[i] = make([]float32, dim, dim)
            var sum float32 = 0.0
            for j := 0; j < dim; j++ {
                embeddingList[i][j] = rand.Float32()
                sum += embeddingList[i][j] * embeddingList[i][j]
            }
            sum = float32(math.Sqrt(float64(sum)))
            for j := 0; j < dim; j++ {
                embeddingList[i][j] = embeddingList[i][j] / sum
            }
            wg.Done()
        }(i)
    }
    wg.Wait()

    return idList, randomList, embeddingList
}

func printResult(sRet *client.SearchResult) {
    idx := make([]int64, 0, sRet.ResultCount)
    randoms := make([]float64, 0, sRet.ResultCount)
    scores := make([]float32, 0, sRet.ResultCount)

    var randCol *entity.ColumnDouble
    var idxCol *entity.ColumnInt64
    for _, field := range sRet.Fields {
        if field.Name() == randomCol {
            c, ok := field.(*entity.ColumnDouble)
            if ok {
                randCol = c
            }
        }
        if field.Name() == idCol {
            c, ok := field.(*entity.ColumnInt64)
            if ok {
                idxCol = c
            }
        }
    }
    for i := 0; i < sRet.ResultCount; i++ {
        val, err := randCol.ValueByIdx(i)
        if err != nil {
            log.Fatal(err)
        }
        randoms = append(randoms, val)
        scores = append(scores, sRet.Scores[i])

        idxVal, err := idxCol.ValueByIdx(i)
        if err != nil {
            log.Fatal(err)
        }
        idx = append(idx, idxVal)
    }
    fmt.Printf("idx: %v\nrandoms: %v\nscores: %v\n", idx, randoms, scores)
}
